import torch
from pytorch_wavelets import DWTForward, DWTInverse

# 创建测试输入
x = torch.randn(4, 32, 480, 640)

# 初始化DWT和IDWT
dwt = DWTForward(J=1, mode='zero', wave='haar')
idwt = DWTInverse(mode='zero', wave='haar')

# 正向小波变换
x_low, x_high = dwt(x)
print(f"x_low shape: {x_low.shape}")    # Expected: [4, 32, 240, 320]
print(f"x_high type: {type(x_high)}")  # Expected: list
print(f"Number of subbands: {len(x_high)}")  # Expected: 3

# 逆向小波变换
x_recon = idwt(x_low, x_high)
print(f"x_recon shape: {x_recon.shape}")  # Expected: [4, 32, 480, 640]
